import tweepy
from datetime import datetime, timedelta
from .base import BaseCrawler

class TwitterCrawler(BaseCrawler):
    """Twitter爬虫实现"""
    
    def __init__(self, task):
        super().__init__(task)
        self.is_running = False
        
        # 从任务参数中获取配置
        self.keywords = self.task.parameters.get('keywords', [])
        self.start_date = self.task.parameters.get('start_date')
        self.end_date = self.task.parameters.get('end_date')
        self.max_tweets = self.task.parameters.get('max_tweets', 100)
        
        # 设置认证信息
        self.credentials = self.task.data_source.credentials
        self.api = self._setup_api()
    
    def _setup_api(self):
        """设置Twitter API客户端"""
        try:
            auth = tweepy.OAuthHandler(
                self.credentials.get('consumer_key'),
                self.credentials.get('consumer_secret')
            )
            auth.set_access_token(
                self.credentials.get('access_token'),
                self.credentials.get('access_token_secret')
            )
            return tweepy.API(auth, wait_on_rate_limit=True)
        except Exception as e:
            self.update_progress(0, f"Twitter API设置失败: {str(e)}")
            return None
    
    def validate_credentials(self):
        """验证Twitter API凭证"""
        try:
            if not self.api:
                return False
            self.api.verify_credentials()
            return True
        except Exception:
            return False
    
    def search_tweets(self, keyword):
        """搜索推文"""
        try:
            query = f"{keyword} -filter:retweets"
            tweets = tweepy.Cursor(
                self.api.search_tweets,
                q=query,
                lang="zh",  # 可以根据需要修改语言
                tweet_mode="extended"
            ).items(self.max_tweets)
            
            return tweets
        except Exception as e:
            self.update_progress(
                self.task.progress,
                f"搜索推文时出错: {str(e)}"
            )
            return []
    
    def parse_tweet(self, tweet):
        """解析推文数据"""
        return {
            'id': tweet.id_str,
            'text': tweet.full_text,
            'created_at': tweet.created_at.isoformat(),
            'retweet_count': tweet.retweet_count,
            'favorite_count': tweet.favorite_count,
            'lang': tweet.lang,
            'user': {
                'id': tweet.user.id_str,
                'screen_name': tweet.user.screen_name,
                'name': tweet.user.name,
                'followers_count': tweet.user.followers_count,
                'friends_count': tweet.user.friends_count,
                'verified': tweet.user.verified
            },
            'hashtags': [tag['text'] for tag in tweet.entities.get('hashtags', [])],
            'mentions': [
                {
                    'screen_name': mention['screen_name'],
                    'id': mention['id_str']
                }
                for mention in tweet.entities.get('user_mentions', [])
            ]
        }
    
    def start(self):
        """开始爬取数据"""
        if not self.validate_credentials():
            self.complete_task(success=False)
            self.update_progress(0, "Twitter API验证失败")
            return
        
        self.is_running = True
        total_processed = 0
        
        try:
            for keyword in self.keywords:
                if not self.is_running:
                    break
                
                tweets = self.search_tweets(keyword)
                for tweet in tweets:
                    if not self.is_running:
                        break
                    
                    parsed_tweet = self.parse_tweet(tweet)
                    self.save_data(
                        data=parsed_tweet,
                        metadata={
                            'keyword': keyword,
                            'crawled_at': datetime.now().isoformat()
                        }
                    )
                    
                    total_processed += 1
                    progress = min(
                        int((total_processed / (len(self.keywords) * self.max_tweets)) * 100),
                        99
                    )
                    self.update_progress(progress)
            
            if self.is_running:
                self.complete_task(success=True)
                self.update_progress(100)
        except Exception as e:
            self.update_progress(
                self.task.progress,
                f"爬取过程中出错: {str(e)}"
            )
            self.complete_task(success=False)
        finally:
            self.is_running = False
    
    def stop(self):
        """停止爬取数据"""
        self.is_running = False 